#!/usr/bin/python3
# coding=UTF-8

import rospy
from std_msgs.msg import Int32
from sensor_msgs.msg import LaserScan
from sensor_msgs.msg import Image
from detection_msgs.msg import Room
import robomaster
from robomaster import robot
from robomaster import camera
from cv_bridge import CvBridge
import math
import numpy as np

# 定义常量
SCAN_RANGE_FRONT = (0, 3)
SCAN_RANGE_LEFT = (88, 92)
SCAN_RANGE_BACK = (178, 182)
SCAN_RANGE_RIGHT = (268, 272)
# 定义标志
is_final = True
# 监听
def lisener(n):
    sum_range = 0
    count = 0
    # 获取雷达数据
    # data = rospy.wait_for_message("/scan", LaserScan, timeout=None)
    try:
        data = rospy.wait_for_message("/scan", LaserScan, timeout=5) 
    except rospy.exceptions.ROSException as e:
        rospy.logerr("Timeout waiting for LaserScan message: %s", e)
        data=None
    # 模式选择
    if n == 0:
        start = SCAN_RANGE_FRONT[0]
        final = SCAN_RANGE_FRONT[1]
    elif n == 1:
        start = SCAN_RANGE_LEFT[0]
        final = SCAN_RANGE_LEFT[1]
    elif n == 2:
        start = SCAN_RANGE_BACK[0]
        final = SCAN_RANGE_BACK[1]
    elif n == 3:
        start = SCAN_RANGE_RIGHT[0]
        final = SCAN_RANGE_RIGHT[1]  
    # 循环选定的索引范围，并跳过inf值
    if data is not None:
        for i in range(start, final):
            if not rospy.is_shutdown():
                if data.ranges[i] != float('inf'):
                    sum_range += data.ranges[i]
                    count += 1
                    rospy.loginfo("scan data:",sum_range/count)
    # 空集判断
    if count:
        return sum_range/count
    else:
        return
# 定义检测角度函数
def check_angle(mode):
    # 起点检测：
    if mode == 0:
        # 迭代三次
        for i in range(3):  
            #scan_data = rospy.wait_for_message("/scan", LaserScan, timeout=None)
            try:
                scan_data = rospy.wait_for_message("/scan", LaserScan, timeout=5) 
            except rospy.exceptions.ROSException as e:
                rospy.logerr("Timeout waiting for LaserScan message: %s", e)
                scan_data=None
            if scan_data is not None:
                ranges = scan_data.ranges
                back_ranges = ranges[:int(len(ranges) * 225 / 360)] + ranges[int(len(ranges) * 135 / 360):]
                wall_indices = np.where(np.array(back_ranges) < 1.5)
                wall_points = np.array([[scan_data.angle_min + i * scan_data.angle_increment, back_ranges[i]] for i in wall_indices[0]])
                if len(wall_points) > 2:
                    slope = np.polyfit(wall_points[:, 0], wall_points[:, 1], 1)
                    wall_angle = math.degrees(np.arctan(slope[0]))
            if wall_angle is not None:
                wall_angle =round(wall_angle * 2,1)
                rospy.loginfo(f"wall_angle: {wall_angle}")
                if wall_angle >= 2:
                    turn_left(wall_angle)
                    rospy.logwarn("Turn left!")
                if wall_angle <= -2:
                    turn_right(-wall_angle)
                    rospy.logwarn("Turn right!")
    # 充电区检测
    elif mode == 1:
        # 迭代三次
        for i in range(3): 
            try:
                scan_data = rospy.wait_for_message("/scan", LaserScan, timeout=5) 
            except rospy.exceptions.ROSException as e:
                rospy.logerr("Timeout waiting for LaserScan message: %s", e)
                scan_data=None
            if scan_data is not None:
                scan_data = rospy.wait_for_message("/scan", LaserScan, timeout=None)
                ranges = scan_data.ranges
                front_ranges = ranges[:int(len(ranges) * 45 / 360)] + ranges[int(len(ranges) * 315 / 360):]
                wall_indices = np.where(np.array(front_ranges) < 1.5)
                wall_points = np.array([[scan_data.angle_min + i * scan_data.angle_increment, front_ranges[i]] for i in wall_indices[0]])
                if len(wall_points) > 2:
                    slope = np.polyfit(wall_points[:, 0], wall_points[:, 1], 1)
                    wall_angle = math.degrees(np.arctan(slope[0]))
        if wall_angle is not None: 
            wall_angle = round(wall_angle * 2,1)
            rospy.loginfo(f"wall_angle: {wall_angle}")
            if wall_angle >= 2:
                turn_left(wall_angle)
                rospy.logwarn("Turn left!")
            elif wall_angle <= -2:
                turn_right(-wall_angle)
                rospy.logwarn("Turn right!")
            else:
                rospy.loginfo("Error free!")
# 定义检查距离函数
def check_distance(mode,target_front_distance,target_left_distance,target_back_distance,target_right_distance):
    # 起终点检测
    if mode == 0:
        # 后向检测
        avg_back = round(lisener(2),3)
        rospy.loginfo("point avg_back : ", avg_back)
        if avg_back != 0:
            error_back = round(target_back_distance - avg_back,3)
            rospy.loginfo("point error_back : ", error_back)
            if abs(error_back) > 0.05:
                ep_chassis.move(x=error_back, y=0, z=0, xy_speed=1.3).wait_for_completed()
        # 右向检测
        avg_right = round(lisener(3),3)
        rospy.loginfo("point avg_right : ", avg_right)
        if avg_right != 0:
            error_right = round(avg_right - target_right_distance,3)
            rospy.loginfo("point error_right : ", error_right)
            if abs(error_right) > 0.05:
                ep_chassis.move(x=0, y=error_right, z=0, xy_speed=1.3).wait_for_completed()
    # 充电区检测
    elif mode == 1:
        # 前向检测
        avg_front = round(lisener(0),3)
        rospy.loginfo("area avg_front : ", avg_front)
        if avg_front != 0:
            error_front = round(avg_front - target_front_distance,3)
            rospy.loginfo("area error_front : ", error_front)
            if abs(error_front) > 0.05:
                ep_chassis.move(x=error_front, y=0, z=0, xy_speed=1.3).wait_for_completed()
        # 左向检测
        avg_left = round(lisener(1),3)
        rospy.loginfo("area avg_left : ", avg_left)
        if avg_left != 0:
            error_left = round(target_left_distance - avg_left,3)
            rospy.loginfo("area error_left : ", error_left)
            if abs(error_left) > 0.05:
                ep_chassis.move(x=0, y=error_left, z=0, xy_speed=1.3).wait_for_completed()
    # 中间检测
    elif mode == 2:
        avg_front = round(lisener(0),3)
        rospy.loginfo("halfway avg_front : ", avg_front)
        if avg_front != 0:
            error_front = round(avg_front - target_front_distance,3)
            rospy.loginfo("halfway error_front : ", error_front)
            if abs(error_front) > 0.1:
                ep_chassis.move(x=error_front, y=0, z=0, xy_speed=1.3).wait_for_completed()
# 定义移动函数
def move_forward(distance):
    ep_chassis.move(x=distance, y=0, z=0, xy_speed=1.7).wait_for_completed()
def move_backward(distance):
    ep_chassis.move(x=-distance, y=0, z=0, xy_speed=1.3).wait_for_completed()
def move_left(distance):
    ep_chassis.move(x=0, y=-distance, z=0, xy_speed=1.7).wait_for_completed()
def move_right(distance):
    ep_chassis.move(x=0, y=distance, z=0, xy_speed=1.7).wait_for_completed()
def turn_left(angle):
    ep_chassis.move(x=0, y=0, z=angle, z_speed=90).wait_for_completed()
def turn_right(angle):
    ep_chassis.move(x=0, y=0, z=-angle, z_speed=90).wait_for_completed()
def stop_state():
    ep_chassis.drive_speed(x=0, y=0, z=0, timeout=5)
    ep_robot.close()
def sleep_state(second):
    rospy.sleep(second)
# 定义相机函数
def camera(ep_camera,row,num):
    i = 0
    while True:
        bridge = CvBridge()
        if i <= 4:
            i += 1
            img = ep_camera.read_cv2_image(strategy="newest")
            if img is not None:
                room_msg=Room()
                room_msg.Img = bridge.cv2_to_imgmsg(img)
                room_msg.Row=row
                room_msg.Num=num
                room_pub.publish(room_msg)
                rospy.loginfo("current image:",i)
            else:
                rospy.logerr("Error!")
        else:
            break
# 全任务
def move(ep_camera):
    # 等待起飞指令
    while True:
        takeoff = rospy.get_param("takeoff_status",0)
        if takeoff == 1:
          break
    rospy.loginfo("Take off from the ground!")
    # 检查起点位置
    check_distance(0,0,0,1,1)
    # 第一排
    row=1
    num=9
    distance = [2,1.5,1,1]
    distance_check = [6.1,4.6,3.6,2.6]
    for i in [0,1,2,3]:
        # 移动
        move_forward(distance[i])
        # 中途检测
        check_distance(2,distance_check[i],0,0,0)
        # 捕获照片
        rospy.loginfo(f"{i+1} Start!")
        ep_camera.start_video_stream(display=False)
        camera(ep_camera,row,num)
        ep_camera.stop_video_stream()
        rospy.loginfo(f"{i+1} Publish!")
        # 防止图像传输不及时
        sleep_time = i + 5
        sleep_state(sleep_time)
        rospy.logwarn(f"{i+1} Done!")
        num+=1
    # 充电区
    move_forward(1.5)
    move_left(2.2)
    # 检查充电区位置
    check_distance(1,1.1,1.4,0,0)
    # 检查充电区姿态
    # check_angle(1)
    # 等待对接指令
    if is_final:
        while True:
            dock = rospy.get_param("dock_status",0)
            if dock == 1:
                break
        rospy.loginfo("Air-to-ground docking!")
        while True:
            takeoff = rospy.get_param("takeoff_status",0)
            if takeoff == 1:
                break
        rospy.loginfo("Take off from the charging area!")
    # 第二排
    distance = [1.3,0.9,1.3,1.3]
    distance_check = [2.6,3.6,4.6,6.1]
    row=2
    num=12
    for i in [0,1,2,3]:
        # 移动
        move_backward(distance[i])
        # 中途检测
        check_distance(2,distance_check[i],0,0,0)
        # 捕获照片
        rospy.loginfo(f"{i+5} Start!")
        ep_camera.start_video_stream(display=False)
        camera(ep_camera,row,num)
        ep_camera.stop_video_stream()
        rospy.loginfo(f"{i+5} Publish!")
        # 防止图像传输不及时
        sleep_time = i + 9
        sleep_state(sleep_time)
        rospy.logwarn(f"{i+5} Done!")
        num-=1
    # 返程
    move_backward(2)
    move_right(2.4)
    # 检查终点位置
    check_distance(0,0,0,0.6,1.2)
# 主函数
if __name__ == '__main__':
    # 初始化节点
    rospy.init_node('ztcar_move')
    # 初始化机器人
    ep_robot = robot.Robot()
    ep_robot.initialize(conn_type="sta")
    ep_chassis = ep_robot.chassis
    ep_camera = ep_robot.camera
    # 初始化参数
    rospy.set_param("takeoff_status",0)
    rospy.set_param("dock_status",0)
    # 发布话题
    #img_pub = rospy.Publisher('car_img',Image,queue_size=10)
    room_pub = rospy.Publisher('car_img',Room,queue_size=10)
    # 运动
    move(ep_camera)
    # 停止运动
    stop_state()
